#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
JPG照片传感器坏点和镜头脏污检测工具
功能：
1. 检测传感器坏点（死点/热点）
2. 检测镜头脏污
"""
import os
import sys
from PIL import Image, ImageDraw, ImageEnhance
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime

class PhotoIssueDetector:
    def __init__(self, sensitivity=2.5, blur_threshold=100, dirty_threshold=0.1):
        """
        初始化检测器
        :param sensitivity: 坏点检测灵敏度 (标准差倍数)
        :param blur_threshold: 模糊检测阈值
        :param dirty_threshold: 脏污检测阈值
        """
        self.sensitivity = sensitivity
        self.blur_threshold = blur_threshold
        self.dirty_threshold = dirty_threshold
        # 创建结果目录
        self.result_dir = 'detection_results'
        os.makedirs(self.result_dir, exist_ok=True)

    def load_image(self, image_path):
        """加载图像"""
        try:
            img = Image.open(image_path)
            return img
        except Exception as e:
            print(f"无法加载图像 {image_path}: {e}")
            return None

    def detect_hot_dead_pixels(self, img):
        """检测热点和死点（优化版）"""
        # 转换为RGB模式
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # 转换为numpy数组
        img_array = np.array(img)
        height, width, _ = img_array.shape

        # 应用更强的高斯模糊减少噪声
        from scipy import ndimage
        blurred_img = ndimage.gaussian_filter(img_array, sigma=1.5)

        # 使用高斯模糊计算局部统计特征 (更内存高效)
        # 计算局部均值 (使用5x5高斯核)
        local_means = ndimage.gaussian_filter(blurred_img, sigma=2, mode='nearest')

        # 计算局部标准差
        # 方法: std = sqrt(E[X^2] - (E[X])^2)
        local_squared_means = ndimage.gaussian_filter(blurred_img**2, sigma=2, mode='nearest')
        local_vars = local_squared_means - local_means**2
        local_stds = np.sqrt(np.maximum(local_vars, 0))  # 避免负数开方

        # 定义局部阈值 (提高阈值以减少热点数量)
        hot_threshold_factor = 2.5  # 提高局部标准差倍数，降低检测灵敏度
        hot_absolute_threshold = 210  # 提高绝对亮度阈值

        # 检测热点 (局部异常亮的点)
        hot_mask = (blurred_img > local_means + hot_threshold_factor * local_stds) & (blurred_img > hot_absolute_threshold)
        hot_mask = np.any(hot_mask, axis=2)  # 任一通道满足条件

        # 应用更强的形态学操作去除孤立噪点
        hot_mask = ndimage.binary_opening(hot_mask, structure=np.ones((7, 7)))

        # 死点检测保持不变
        mean_r = np.mean(blurred_img[:, :, 0])
        mean_g = np.mean(blurred_img[:, :, 1])
        mean_b = np.mean(blurred_img[:, :, 2])
        std_r = np.std(blurred_img[:, :, 0])
        std_g = np.std(blurred_img[:, :, 1])
        std_b = np.std(blurred_img[:, :, 2])

        dead_threshold_r = mean_r - self.sensitivity * 1.5 * std_r
        dead_threshold_g = mean_g - self.sensitivity * 1.5 * std_g
        dead_threshold_b = mean_b - self.sensitivity * 1.5 * std_b
        dead_mask = (blurred_img[:, :, 0] < dead_threshold_r) & \
                    (blurred_img[:, :, 0] < 50) | \
                    (blurred_img[:, :, 1] < dead_threshold_g) & \
                    (blurred_img[:, :, 1] < 50) | \
                    (blurred_img[:, :, 2] < dead_threshold_b) & \
                    (blurred_img[:, :, 2] < 50)

        # 应用更强的形态学操作去除孤立噪点
        hot_mask = ndimage.binary_opening(hot_mask, structure=np.ones((5, 5)))
        dead_mask = ndimage.binary_opening(dead_mask, structure=np.ones((3, 3)))

        # 转换掩码为坐标列表
        hot_pixels = list(zip(*np.where(hot_mask)))[::-1]  # 转换为(x,y)格式
        dead_pixels = list(zip(*np.where(dead_mask)))[::-1]  # 转换为(x,y)格式

        return hot_pixels, dead_pixels

    def detect_lens_dirt(self, img):
        """检测镜头脏污（优化版）"""
        # 转换为灰度图
        gray = img.convert('L')
        
        # 转换为numpy数组
        gray_array = np.array(gray)
        height, width = gray_array.shape
        
        # 应用高斯模糊减少噪声
        from scipy import ndimage
        blurred_gray = ndimage.gaussian_filter(gray_array, sigma=3)
        
        # 计算局部对比度 (使用13x13邻域) - 使用滑动窗口优化
        kernel_size = 13
        pad = kernel_size // 2
        padded_gray = np.pad(blurred_gray, pad, mode='edge')
        
        # 创建滑动窗口视图
        from numpy.lib.stride_tricks import sliding_window_view
        windows = sliding_window_view(padded_gray, (kernel_size, kernel_size))
        
        # 计算每个窗口的均值和标准差
        local_means = np.mean(windows, axis=(2, 3))
        local_stds = np.std(windows, axis=(2, 3))
        
        # 计算对比度图
        contrast_map = local_stds / (local_means + 1e-6)  # 避免除以零
        
        # 计算全局对比度特征
        global_mean_contrast = np.mean(contrast_map)
        global_std_contrast = np.std(contrast_map)
        
        # 设置自适应阈值 - 降低阈值提高脏污检测灵敏度
        adaptive_threshold = global_mean_contrast - 1.5 * global_std_contrast
        adaptive_threshold = max(0.05, adaptive_threshold)  # 降低最小阈值
        
        # 找出低对比度区域
        low_contrast_regions = np.where(contrast_map < adaptive_threshold)
        
        # 标记脏污区域
        dirty_regions = list(zip(low_contrast_regions[1], low_contrast_regions[0]))  # (x,y)格式
        
        # 过滤掉可能是均匀区域而不是脏污的区域
        # 计算图像的梯度
        gradient = ndimage.sobel(blurred_gray)
        gradient_threshold = np.mean(gradient) * 0.95  # 更高阈值以减少误判
        
        # 使用向量化操作过滤梯度
        if len(dirty_regions) > 0:
            # 将脏污区域转换为numpy数组
            dirty_coords = np.array(dirty_regions)
            x_coords, y_coords = dirty_coords[:, 0], dirty_coords[:, 1]
            
            # 检查梯度阈值
            gradient_mask = gradient[y_coords, x_coords] > gradient_threshold
            filtered_dirty_regions = dirty_coords[gradient_mask]
            filtered_dirty_regions = list(map(tuple, filtered_dirty_regions))
        else:
            filtered_dirty_regions = []
        
        # 应用形态学操作进一步去除噪声
        if len(filtered_dirty_regions) > 0:
            # 创建掩码
            mask = np.zeros((height, width), dtype=bool)
            if filtered_dirty_regions:
                x_coords, y_coords = zip(*filtered_dirty_regions)
                mask[y_coords, x_coords] = True
            
            # 应用闭操作填充小漏洞
            closed_mask = ndimage.binary_closing(mask, structure=np.ones((7, 7)))
            
            # 应用开操作去除小的噪点
            opened_mask = ndimage.binary_opening(closed_mask, structure=np.ones((5, 5)))
            
            # 找出连通区域
            labeled_mask, num_features = ndimage.label(opened_mask)
            
            # 只保留面积和形状符合脏污特征的区域
            min_area = 500  # 减小最小面积阈值，检测更小脏污
            max_aspect_ratio = 5.0  # 放宽长宽比限制
            border_margin = 10  # 减小边缘边距，允许检测边缘脏污
            filtered_regions = np.zeros_like(labeled_mask, dtype=bool)
            
            # 使用向量化操作计算区域属性
            for i in range(1, num_features + 1):
                region = labeled_mask == i
                region_area = np.sum(region)
                
                # 过滤小区域
                if region_area < min_area:
                    continue
                
                # 计算区域的边界框
                coords = np.where(region)
                y_min, y_max = np.min(coords[0]), np.max(coords[0])
                x_min, x_max = np.min(coords[1]), np.max(coords[1])
                
                # 检查是否靠近图像边缘
                near_edge = (x_min < border_margin or x_max > width - border_margin or
                            y_min < border_margin or y_max > height - border_margin)
                if near_edge:
                    continue  # 过滤边缘区域
                
                # 计算长宽比
                height_region = y_max - y_min + 1
                width_region = x_max - x_min + 1
                aspect_ratio = max(height_region, width_region) / min(height_region, width_region + 1e-6)
                
                # 过滤过长的区域
                if aspect_ratio > max_aspect_ratio:
                    continue
                
                # 计算区域的紧凑度 (圆形度)
                # 使用膨胀操作计算周长
                dilated = ndimage.binary_dilation(region, structure=np.ones((3,3)))
                perimeter = np.sum(dilated) - np.sum(region)
                if perimeter == 0:
                    continue
                compactness = 4 * np.pi * region_area / (perimeter ** 2)
                
                # 保留更接近圆形的区域 (脏污通常是圆形或类圆形)
                if compactness > 0.3:
                    filtered_regions |= region
            
            # 更新过滤后的脏污区域
            y_coords, x_coords = np.where(filtered_regions)
            filtered_dirty_regions = list(zip(x_coords, y_coords))
        
        # 重新计算脏污百分比
        filtered_dirty_percentage = len(filtered_dirty_regions) / (height * width)
        
        return filtered_dirty_regions, filtered_dirty_percentage

    def visualize_results(self, img, hot_pixels, dead_pixels, dirty_regions, image_path):
        """可视化检测结果"""
        # 创建图像副本
        result_img = img.copy()
        draw = ImageDraw.Draw(result_img)

        # 标记热点（红色）
        for x, y in hot_pixels:
            draw.ellipse([(x-2, y-2), (x+2, y+2)], fill='red', outline='red')

        # 标记死点（蓝色）
        for x, y in dead_pixels:
            draw.ellipse([(x-2, y-2), (x+2, y+2)], fill='blue', outline='blue')

        # 标记脏污区域（黄色半透明区域）
        # 创建脏污区域的掩码
        width, height = img.size
        dirt_mask = Image.new('RGBA', (width, height), (0, 0, 0, 0))
        dirt_draw = ImageDraw.Draw(dirt_mask)
        
        # 将脏污点集转换为 numpy 数组以便处理
        dirty_points = np.array(dirty_regions)
        
        if len(dirty_points) > 0:
            # 增加标记数量上限，根据图像大小动态调整
            max_points = min(len(dirty_points), width * height // 100)  # 最多标记图像像素的1%
            sampled_points = dirty_points[:max_points]
            
            # 绘制半透明黄色区域
            for x, y in sampled_points:
                dirt_draw.ellipse([(x-1, y-1), (x+1, y+1)], fill=(255, 255, 0, 100))
        
        # 将脏污掩码合并到结果图像
        result_img = Image.alpha_composite(result_img.convert('RGBA'), dirt_mask)
        result_img = result_img.convert('RGB')  # 转换回RGB模式

        # 保存结果图像
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = os.path.basename(image_path)
        result_filename = f"{self.result_dir}/{filename}_result_{timestamp}.jpg"
        result_img.save(result_filename)

        print(f"检测结果已保存至: {result_filename}")
        return result_filename

    def generate_report(self, image_path, hot_pixels, dead_pixels, dirty_percentage):
        """生成检测报告"""
        report = {
            'image_path': image_path,
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'hot_pixels_count': len(hot_pixels),
            'dead_pixels_count': len(dead_pixels),
            'dirty_percentage': round(dirty_percentage * 100, 2),
            'hot_pixels_detected': len(hot_pixels) > 0,
            'dead_pixels_detected': len(dead_pixels) > 0,
            'dirt_detected': dirty_percentage > self.dirty_threshold
        }

        # 打印报告
        print("=== 检测报告 ===")
        print(f"图像路径: {report['image_path']}")
        print(f"检测时间: {report['timestamp']}")
        print(f"热点数量: {report['hot_pixels_count']}")
        print(f"死点数量: {report['dead_pixels_count']}")
        print(f"脏污百分比: {report['dirty_percentage']}%")
        print(f"是否检测到热点: {'是' if report['hot_pixels_detected'] else '否'}")
        print(f"是否检测到死点: {'是' if report['dead_pixels_detected'] else '否'}")
        print(f"是否检测到脏污: {'是' if report['dirt_detected'] else '否'}")
        print("================")

        return report

    def detect(self, image_path):
        """执行完整检测流程"""
        print(f"开始检测图像: {image_path}")
        img = self.load_image(image_path)
        if img is None:
            return None

        # 检测坏点
        hot_pixels, dead_pixels = self.detect_hot_dead_pixels(img)

        # 检测脏污
        dirty_regions, dirty_percentage = self.detect_lens_dirt(img)

        # 可视化结果
        result_img_path = self.visualize_results(img, hot_pixels, dead_pixels, dirty_regions, image_path)

        # 生成报告
        report = self.generate_report(image_path, hot_pixels, dead_pixels, dirty_percentage)

        return report

def main():
    if len(sys.argv) < 2:
        print("用法: python detect_photo_issues.py <图像路径1> [<图像路径2> ...]")
        return

    detector = PhotoIssueDetector(sensitivity=2.5, blur_threshold=80, dirty_threshold=0.1)

    for image_path in sys.argv[1:]:
        if not os.path.exists(image_path):
            print(f"图像文件不存在: {image_path}")
            continue
        detector.detect(image_path)

if __name__ == '__main__':
    main()